#!/usr/bin/env python3
"""
Demonstration of covariance matrix validation
"""

import torch
import numpy as np
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")


class CovarianceDemo:
    """Demo class to show covariance matrix validation"""

    def _is_positive_definite(self, matrix, tolerance=1e-8):
        """Check if a matrix is positive definite"""
        try:
            torch.linalg.cholesky(matrix)
            return True
        except RuntimeError:
            eigenvals = torch.linalg.eigvals(matrix).real
            return torch.all(eigenvals > tolerance)

    def _validate_covariance_matrix(self, cov_matrix, class_idx):
        """Validate covariance matrix properties"""
        print(f"\n=== Validating Covariance Matrix for Class {class_idx} ===")

        # Check positive definiteness
        is_pd = self._is_positive_definite(cov_matrix)
        print(f"1. Positive Definite: {is_pd}")

        if not is_pd:
            logging.error(
                f"Covariance matrix for class {class_idx} is not positive definite"
            )
            # In real code, this would raise an exception
            print("   ❌ FAILED: Matrix is not positive definite!")
            return False

        # Check condition number
        try:
            cond_num = torch.linalg.cond(cov_matrix)
            print(f"2. Condition Number: {cond_num:.2e}")

            if cond_num > 1e12:
                logging.warning(f"High condition number: {cond_num}")
                print(
                    "   ⚠️  WARNING: High condition number (numerical instability risk)"
                )
            else:
                print("   ✅ Good condition number")

        except RuntimeError:
            logging.warning(f"Could not compute condition number for class {class_idx}")
            print("   ⚠️  WARNING: Could not compute condition number")

        # Check determinant
        det = torch.linalg.det(cov_matrix)
        print(f"3. Determinant: {det:.2e}")

        if det <= 0:
            logging.error(f"Non-positive determinant: {det}")
            print("   ❌ FAILED: Non-positive determinant!")
            return False
        else:
            print("   ✅ Positive determinant")

        # Check eigenvalues
        eigenvals = torch.linalg.eigvals(cov_matrix).real
        min_eigenval = torch.min(eigenvals)
        max_eigenval = torch.max(eigenvals)

        print(f"4. Eigenvalue Range: [{min_eigenval:.2e}, {max_eigenval:.2e}]")

        if min_eigenval <= 0:
            print("   ❌ FAILED: Non-positive eigenvalues!")
            return False
        else:
            print("   ✅ All eigenvalues positive")

        # Overall success
        logging.info(
            f"Covariance matrix for class {class_idx} validation passed (det={det:.2e})"
        )
        print("   🎉 VALIDATION PASSED!")
        return True


def create_test_matrices():
    """Create different types of matrices for testing"""

    # Good matrix: well-conditioned, positive definite
    good_matrix = torch.tensor([[2.0, 0.5], [0.5, 3.0]], dtype=torch.float64)

    # Bad matrix: not positive definite
    bad_matrix = torch.tensor(
        [[1.0, 2.0], [2.0, 1.0]],  # This will have negative eigenvalues
        dtype=torch.float64,
    )

    # Singular matrix: determinant = 0
    singular_matrix = torch.tensor(
        [[1.0, 2.0], [2.0, 4.0]], dtype=torch.float64  # Second row = 2 * first row
    )

    # Ill-conditioned matrix: high condition number
    ill_conditioned = torch.tensor(
        [[1.0, 0.99999], [0.99999, 1.0]], dtype=torch.float64
    )

    return {
        "good": good_matrix,
        "bad": bad_matrix,
        "singular": singular_matrix,
        "ill_conditioned": ill_conditioned,
    }


def main():
    """Test different matrix types"""
    demo = CovarianceDemo()
    test_matrices = create_test_matrices()

    print("Covariance Matrix Validation Demo")
    print("=" * 60)

    for matrix_type, matrix in test_matrices.items():
        print(f"\nTesting {matrix_type.upper()} matrix:")
        print(f"Matrix:\n{matrix}")

        success = demo._validate_covariance_matrix(matrix, matrix_type)

        if success:
            print("✅ This matrix is suitable for use in multivariate distributions")
        else:
            print("❌ This matrix needs regularization before use")

        print("-" * 50)


if __name__ == "__main__":
    main()
